package com.gjy.kafka.sp.config;

import com.gjy.kafka.sp.entity.MessageEntity;
import org.apache.commons.lang3.StringUtils;
import org.apache.kafka.clients.producer.Partitioner;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.InvalidRecordException;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.utils.Utils;

import java.util.List;
import java.util.Map;

/**
 * 自定义Kafka分区策略
 * 根据业务ID将消息发送到指定分区，确保相同业务ID的消息在同一分区
 *
 * @author gjy
 * @version 1.0
 * @since 2025-09-07 10:38:44
 */
public class BusinessIdPartitioner implements Partitioner {

    /**
     * 计算分区号
     *
     * @param topic      主题名称
     * @param key        消息键
     * @param keyBytes   消息键的字节数组
     * @param value      消息值
     * @param valueBytes 消息值的字节数组
     * @param cluster    Kafka集群信息
     * @return 分区号
     */
    @Override
    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
        // 获取主题的所有分区
        List<PartitionInfo> partitions = cluster.partitionsForTopic(topic);
        int numPartitions = partitions.size();

        // 如果消息值不是MessageEntity类型，抛出异常
        if (!(value instanceof MessageEntity)) {
            throw new InvalidRecordException("消息必须是MessageEntity类型");
        }

        MessageEntity message = (MessageEntity) value;
        String businessId = message.getBusinessId();

        // 如果业务ID为空，使用默认分区策略
        if (StringUtils.isBlank(businessId)) {
            if (keyBytes == null) {
                // 使用随机分区
                return Utils.toPositive(Utils.murmur2(valueBytes)) % numPartitions;
            } else {
                // 使用key计算分区
                return Utils.toPositive(Utils.murmur2(keyBytes)) % numPartitions;
            }
        }

        // 根据业务ID计算分区，确保相同业务ID的消息在同一分区
        return Math.abs(businessId.hashCode()) % numPartitions;
    }

    /**
     * 关闭分区器
     */
    @Override
    public void close() {
        // 关闭资源（如果有的话）
    }

    /**
     * 配置分区器
     *
     * @param configs 配置参数
     */
    @Override
    public void configure(Map<String, ?> configs) {
        // 读取配置参数（如果有的话）
    }

}
